import requests
import time
import os
import json
import webbrowser
from requests.exceptions import ChunkedEncodingError

try:
    import msvcrt  # Windows-specific
except ImportError:
    msvcrt = None

TOKEN_FILE = 'token.txt'
API_URL = 'https://xz.klszkj.com:9000/botapi/api/llm/v1/completion'
DEFAULT_LOOPS = 20
INPUT_TIMEOUT = 5  # seconds

# 定义可自定义的问题列表
QUERIES = [
    'Java是什么？',
    'Python的优缺点是什么？',
    'Nginx是什么？',
    'Tomcat是什么？',
    'c++是什么？',
    'nodejs是什么？',
    'php是什么？',
    'java的优缺点是什么？',
    'php的优缺点是什么？',
    '如何开始学习java？',
    '如何开始学习C++？'
]


def read_token():
    if os.path.exists(TOKEN_FILE):
        with open(TOKEN_FILE, 'r') as f:
            return f.read().strip()
    return None


def save_token(token):
    with open(TOKEN_FILE, 'w') as f:
        f.write(token)


def prompt_token():
    token = input('输入昆仑小智中的Token: ').strip()
    save_token(token)
    return token


def input_with_timeout(prompt, timeout, default):
    if msvcrt:
        print(prompt, end='', flush=True)
        start = time.time()
        input_str = ''
        while True:
            if msvcrt.kbhit():
                ch = msvcrt.getwche()
                # Enter keys
                if ch in ('\r', '\n'):
                    print()
                    break
                input_str += ch
            if time.time() - start > timeout:
                print()  # newline
                return default
        return int(input_str) if input_str.isdigit() else default
    else:
        # Fallback for non-Windows
        from threading import Thread
        from queue import Queue, Empty
        q = Queue()
        def _get_input():
            try:
                q.put(input(prompt))
            except Exception:
                q.put(None)
        t = Thread(target=_get_input)
        t.daemon = True
        t.start()
        try:
            val = q.get(timeout=timeout)
            return int(val) if val and val.isdigit() else default
        except Empty:
            print()  # newline
            return default


def send_request(session, headers, payload):
    '''发送请求，如果401/403则提示用户重新输入Token并重试一次'''
    try:
        resp = session.post(API_URL, headers=headers, json=payload, stream=True, timeout=10)
    except requests.exceptions.RequestException as e:
        print('Request error:', e)
        return None

    if resp.status_code in (401, 403):
        print('Token无效或已过期。请输入新Token。')
        new_token = prompt_token()
        headers['Authorization'] = f'Bearer {new_token}'
        headers['token'] = new_token
        save_token(new_token)
        try:
            resp = session.post(API_URL, headers=headers, json=payload, stream=True, timeout=10)
        except requests.exceptions.RequestException as e:
            print('重试请求失败: ', e)
            return None
    return resp


def main():
    loops = input_with_timeout(
        f'在 {INPUT_TIMEOUT} 秒内输入请求次数 (默认 {DEFAULT_LOOPS})次',
        INPUT_TIMEOUT,
        DEFAULT_LOOPS
    )
    print(f'将发送 {loops} 次请求。')

    token = read_token() or prompt_token()
    headers = {
        'Accept': 'text/event-stream',
        'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
        'Content-Type': 'application/json',
        'Origin': 'https://xz.klszkj.com:9000',
        'Authorization': f'Bearer {token}',
        'token': token
    }

    session = requests.Session()


    for i in range(1, loops + 1):
        question = QUERIES[(i - 1) % len(QUERIES)]
        print(f'--- Request {i}/{loops}: {question} ---')
        payload = {
            'appId': '0OI8OHvGqeHjtMavFBGE2k',
            'convId': '',
            'uuid': None,
            'files': [],
            'query': question,
            'welcome': '',
            'deep': True,
            'inputs': {}
        }

        response = send_request(session, headers, payload)
        if not response:
            continue
        print('响应状态: ', response.status_code)

        answers = []
        try:
            for line in response.iter_lines(decode_unicode=True):
                if not line or not line.startswith('data:'):
                    continue
                raw = line[len('data:'):]
                try:
                    msg = json.loads(raw)
                except json.JSONDecodeError:
                    continue
                data = msg.get('data', {})
                outputs = data.get('outputs', {})
                ans = outputs.get('answer')
                if ans:
                    answers.append(ans)
                    print(ans)
        except ChunkedEncodingError:
            print('警告: 响应提前结束，继续...')
        except Exception as e:
            print('读取响应时出错: ', e)

        if answers:
            print('完成回答: ', answers[-1])
        time.sleep(0.1)

if __name__ == '__main__':
    webbrowser.open("https://xz.klszkj.com:9000/botapi/chat-v2/0OI8OHvGqeHjtMavFBGE2k")
    main()

